#include "BTreeNode.h"

using namespace std;

/**
 * Constructor, set fields within the BTLeafNode that don't change
 */
BTLeafNode::BTLeafNode() {
  pKeyCount = (int *) buffer;
  pbtle = (BTLeafElement *) (buffer + sizeof(int));
  pNextPid = (PageId *) (buffer + sizeof(int) + sizeof(BTLeafElement) * KEYS_PER_LEAF);

  memset(buffer, 0x00, sizeof(buffer));
}


/*
 * Read the content of the node from the page pid in the PageFile pf.
 * @param pid[IN] the PageId to read
 * @param pf[IN] PageFile to read from
 * @return 0 if successful. Return an error code if there is an error.
 */
RC BTLeafNode::read(PageId pid, const PageFile& pf)
{ 
  // Load page
  RC rc;
  if ((rc = pf.read(pid, buffer)))
    return rc;

  // Keep track of the node's pid
  nodePid = pid;

  return 0;
}
    
/*
 * Write the content of the node to the page pid in the PageFile pf.
 * @param pid[IN] the PageId to write to
 * @param pf[IN] PageFile to write to
 * @return 0 if successful. Return an error code if there is an error.
 */
RC BTLeafNode::write(PageId pid, PageFile& pf)
{ 
  return pf.write(pid, buffer);
}

/*
 * Return the number of keys stored in the node.
 * @return the number of keys in the node
 */
int BTLeafNode::getKeyCount()
{
  return *pKeyCount;
}

/*
 * Insert a (key, rid) pair to the node.
 * @param key[IN] the key to insert
 * @param rid[IN] the RecordId to insert
 * @return 0 if successful. Return an error code if the node is full.
 */
RC BTLeafNode::insert(int key, const RecordId& rid)
{ 
  // Check if there is space for the (key, rid) pair
  if (*pKeyCount == KEYS_PER_LEAF)
    return RC_NODE_FULL;

  // Search for where the (key, rid) pair should be inserted
  int eid;
  if (locate(key, eid))
    eid = *pKeyCount;

  // Move all nodes after and including loc back to make room for new pair
  if (eid < *pKeyCount)
    memmove(pbtle + eid + 1, pbtle + eid, (*pKeyCount - eid) * sizeof(BTLeafElement));

  // Write pair
  pbtle[eid].rid = rid;
  pbtle[eid].key = key;

  // Increment key count
  (*pKeyCount)++;

  return 0;
}

/*
 * Insert the (key, rid) pair to the node
 * and split the node half and half with sibling.
 * The first key of the sibling node is returned in siblingKey.
 * @param key[IN] the key to insert.
 * @param rid[IN] the RecordId to insert.
 * @param sibling[IN] the sibling node to split with. This node MUST be EMPTY when this function is called.
 * @param siblingKey[OUT] the first key in the sibling node after split.
 * @return 0 if successful. Return an error code if there is an error.
 */
RC BTLeafNode::insertAndSplit(int key, const RecordId& rid, 
                              BTLeafNode& sibling, int& siblingKey)
{ 
  // Locate where the key should be inserted
  int eid;
  if (locate(key, eid))
    eid = *pKeyCount;

  // split is ceil(keyCount/2)
  int split = (*pKeyCount + 1) / 2;

  // Get number of elements to move to sibling
  int nCopyEls;
  if (eid >= split) 
    nCopyEls = *pKeyCount - split;
  else
    nCopyEls = split;

  // Do Copy
  memcpy(sibling.pbtle, pbtle + (*pKeyCount - nCopyEls), nCopyEls * sizeof(BTLeafElement));

  // Set key counts in sibling and this node
  *(sibling.pKeyCount) = nCopyEls;
  *pKeyCount = *pKeyCount - nCopyEls;

  if (eid >= split)
    // Insert (key, rid) pair into sibling
    sibling.insert(key, rid);
  else
    insert(key, rid);

  // Return sibling key
  siblingKey = sibling.pbtle[0].key;

  // Update next pid pointers
  PageId temppid = *pNextPid;
  *pNextPid = sibling.nodePid;
  *(sibling.pNextPid) = temppid;

  return 0; 
}

/*
 * Find the entry whose key value is larger than or equal to searchKey
 * and output the eid (entry number) whose key value >= searchKey.
 * Remeber that all keys inside a B+tree node should be kept sorted.
 * @param searchKey[IN] the key to search for
 * @param eid[OUT] the entry number that contains a key larger than or equalty to searchKey
 * @return 0 if successful. Return an error code if there is an error.
 */
RC BTLeafNode::locate(int searchKey, int& eid)
{ 
  // Sequentially search for key
  int ceid;
  for (ceid = 0; ceid < *pKeyCount; ceid++) {
    if (pbtle[ceid].key >= searchKey) {
      eid = ceid;
      return 0;
    }
  }

  // Key >= searchKey not found
  return RC_NO_SUCH_RECORD;
}

/*
 * Read the (key, rid) pair from the eid entry.
 * @param eid[IN] the entry number to read the (key, rid) pair from
 * @param key[OUT] the key from the entry
 * @param rid[OUT] the RecordId from the entry
 * @return 0 if successful. Return an error code if there is an error.
 */
RC BTLeafNode::readEntry(int eid, int& key, RecordId& rid)
{ 
  // Check if eid is a valid entry
  if (eid >= *pKeyCount)
    return RC_NO_SUCH_RECORD;

  // Set the key and rid
  key = pbtle[eid].key;
  rid = pbtle[eid].rid;

  return 0; 
}

/*
 * Return the pid of the next slibling node.
 * @return the PageId of the next sibling node 
 */
PageId BTLeafNode::getNextNodePtr()
{ 
  return *pNextPid;
}

/*
 * Set the pid of the next slibling node.
 * @param pid[IN] the PageId of the next sibling node 
 * @return 0 if successful. Return an error code if there is an error.
 */
RC BTLeafNode::setNextNodePtr(PageId pid)
{ 
  *pNextPid = pid;
  return 0;
}

/**
 * Constructor, set fields within the BTNonLeafNode that don't change
 */
BTNonLeafNode::BTNonLeafNode() {
  pKeyCount = (int *) buffer;
  pbtnle = (BTNonLeafElement *) (buffer + sizeof(int));
  pLastPid = (PageId *) (buffer + sizeof(int) + sizeof(BTNonLeafElement) * KEYS_PER_NONLEAF);

  memset(buffer, 0x00, sizeof(buffer));
}

/*
 * Read the content of the node from the page pid in the PageFile pf.
 * @param pid[IN] the PageId to read
 * @param pf[IN] PageFile to read from
 * @return 0 if successful. Return an error code if there is an error.
 */
RC BTNonLeafNode::read(PageId pid, const PageFile& pf)
{
  return pf.read(pid, buffer);
}
    
/*
 * Write the content of the node to the page pid in the PageFile pf.
 * @param pid[IN] the PageId to write to
 * @param pf[IN] PageFile to write to
 * @return 0 if successful. Return an error code if there is an error.
 */
RC BTNonLeafNode::write(PageId pid, PageFile& pf)
{
  return pf.write(pid, buffer);
}
/*
 * Return the number of keys stored in the node.
 * @return the number of keys in the node
 */
int BTNonLeafNode::getKeyCount()
{
  return *pKeyCount;
}

/*
 * Insert a (key, pid) pair to the node.
 * @param key[IN] the key to insert
 * @param pid[IN] the PageId to insert
 * @return 0 if successful. Return an error code if the node is full.
 */
RC BTNonLeafNode::insert(int key, PageId pid, bool atEnd)
{
  // Check if node is full
  if (*pKeyCount == KEYS_PER_NONLEAF)
    return RC_NODE_FULL;

  // Find where to insert the new entry
  int eid;
  for (eid = 0; eid < *pKeyCount; eid++)
    if (pbtnle[eid].key >= key)
      break;

  // Make room for new element if necessary
  if (eid < *pKeyCount) {
    memmove(pbtnle + eid + 1, pbtnle + eid, (*pKeyCount - eid) * sizeof(BTNonLeafElement));

    // Add element
    pbtnle[eid].pid = pid;
    pbtnle[eid].key = key;
  }
  // New element is added to the end
  else {
    // The new element will be a (last pid, key) pair
    // The supplied pid will become the last pid
    if (atEnd) {
      pbtnle[eid].pid = *pLastPid;
      pbtnle[eid].key = key;
      *pLastPid = pid;
    }
    else {
      *pLastPid = pbtnle[*pKeyCount - 1].pid;
      pbtnle[*pKeyCount - 1].pid = pid;
      pbtnle[*pKeyCount - 1].key = key;
    }
  }

  // Increment key count
  (*pKeyCount)++;

  return 0;
}

/*
 * Split the node half and half with sibling.
 * The middle key after the split is returned in midKey.
 * @param key[IN] the key to insert
 * @param pid[IN] the pid to insert
 * @param sibling[IN] the sibling node to split with. This node MUST be empty when this function is called.
 * @param midKey[OUT] the key in the middle after the split. This key should be inserted to the parent node.
 * @return 0 if successful. Return an error code if there is an error.
 */
RC BTNonLeafNode::insertAndSplit(int key, const PageId& pid, BTNonLeafNode& sibling, int& midKey)
{ 
  // Find the split, ceil(keyCount/2)
  int split = (*pKeyCount + 1) / 2;

  // Find where the new key would go
  int eid;
  for (eid = 0; eid < *pKeyCount; eid++)
    if (key <= pbtnle[eid].key)
      break;

  // If eid >= split, the (key, pid) pair will go into the sibling
  // Otherwise it goes into the current node

  // Get number of elements to copy to sibling
  int nCopyEls;
  if (eid >= split) 
    nCopyEls = *pKeyCount - split - 1;
  else
    nCopyEls = split;

  // Do copy to sibling
  memcpy(sibling.pbtnle, pbtnle + (*pKeyCount - nCopyEls), nCopyEls * sizeof(BTNonLeafElement));

  // Set key counts
  *(sibling.pKeyCount) = nCopyEls;
  *pKeyCount = *pKeyCount - nCopyEls;

  // Set new last pid's
  // Sibling takes the current node's last pid
  *(sibling.pLastPid) = *pLastPid;
  
  // If we would put this at the split, then it just becomes the last pid
  // of this node
  if (eid == split) {
    *pLastPid = pid;
  }
  else {  
    // Insert the key
    if (eid > split)
      sibling.insert(key, pid, false);
    else if (eid < split)
      insert(key, pid, false);

    // In the current node, drop the last key and use the pid as the last pid
    *pLastPid = pbtnle[*pKeyCount - 1].pid;
    (*pKeyCount)--;
  
  }

  // Return midKey
  midKey = sibling.pbtnle[0].key;
  
  return 0;
}

/*
 * Given the searchKey, find the child-node pointer to follow and
 * output it in pid.
 * @param searchKey[IN] the searchKey that is being looked up.
 * @param pid[OUT] the pointer to the child node to follow.
 * @return 0 if successful. Return an error code if there is an error.
 */
RC BTNonLeafNode::locateChildPtr(int searchKey, PageId& pid)
{ 
  // Find location of matching key
  int eid;
  for (eid = 0; eid < *pKeyCount; eid++) {
    // If the current key is greater than the searchKey, 
    // follow the PageId pointer
    if (pbtnle[eid].key >= searchKey) {
      pid = pbtnle[eid].pid;
      return 0;
    }
  }

  // End of keys reached, return last pid
  pid = *pLastPid;
  return 0;
}

/*
 * Initialize the root node with (pid1, key, pid2).
 * @param pid1[IN] the first PageId to insert
 * @param key[IN] the key that should be inserted between the two PageIds
 * @param pid2[IN] the PageId to insert behind the key
 * @return 0 if successful. Return an error code if there is an error.
 */
RC BTNonLeafNode::initializeRoot(PageId pid1, int key, PageId pid2)
{ 
  // Set key count to 1
  *pKeyCount = 1;

  // Set first (key, pid) pair
  pbtnle[0].pid = pid1;
  pbtnle[0].key = key;

  // Set last pid as pid2
  *pLastPid = pid2;

  return 0; 
}
